/*
 * Grammar.java
 *
 * Created on July 29, 2006, 7:44 PM
 *
 */

package galronnlp.perceptron;

import galronnlp.cfg.grammar.Rule;
import galronnlp.util.Production;
import galronnlp.util.ProductionGenerator;
import galronnlp.util.Symbol;
import galronnlp.util.SymbolGenerator;
import galronnlp.util.Tuple;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.LinkedList;

/**
 * This is the grammar object for the perceptron parser as described in Collins & Roark 2002
 *
 * @author Daniel A. Galron
 */
public class Grammar {

    /**  
     * The data structure to hold the chains is a Hashtable of Hashtables of 
     * LinkedLists of LinkedLists: The outer Hashtable  is indexed by the last 
     * Symbol in the chain; the inner Hashtable is indexed by the first Symbol 
     * of the chain; the HashSet is a HashSet containing each chain in the inner 
     * Hashtable.
     */
    
    public Hashtable<Symbol, 
           Hashtable<Symbol, 
           HashSet<Chain>>> chains = new Hashtable<Symbol,
                                         Hashtable<Symbol,
                                         HashSet<Chain>>>();
    
    /**
     * The data structure to hold the allowable triples is a Hashtable of HashSets
     * of Tuples; the Hashtable is indexed by the symbol which is the head of the
     * candidate chain to attach to a subtree; each Tuple in the HashSet contains
     * the allowable parent and left sibling respectively.
     */
    public Hashtable<Symbol, HashSet<Tuple>> allowableTriples =
            new Hashtable<Symbol, HashSet<Tuple>>();
    
    // This hashstable contains the number of times a chain has been seen.
    private Hashtable<Chain, Integer> chainCounts = new Hashtable<Chain, Integer>();
    
    /**
     * This Hashset holds the CFG productions of the grammar, from which
     * the chains and allowable triples were constructed.
     */
    private HashSet<Production> productions = new HashSet<Production>();
    
    /** 
     * This Hashset holds the non-terminal symbols
     */
    private HashSet<Symbol> nonTerminals = new HashSet<Symbol>();
    
    /**
     * This Hashset holds the terminal symbols
     */
    private HashSet<Symbol> Terminals = new HashSet<Symbol>();

    /**
     * This Hashtable maps Symbol to Symbol - a lexical item (i.e. a word) to
     * its POS tag
     */
    public Hashtable<Symbol, Symbol> lexicon = new Hashtable<Symbol, Symbol>();
    
    /** 
     * The start symbol
     */
    private Symbol StartSymbol;
    
    /** 
     * The end-of-constituent symbol
     */
    private Symbol EndSymbol;
    
    /**
     * a reference to the parameter vector
     */
    public FeatureVector parameters = new FeatureVector();
    
    /**
     * A reference to a ProductionGenerator object, which assigns a number to
     * a rule to allow for quick comparison
     */
    public ProductionGenerator prodGen = new ProductionGenerator();
    
    /**
     * A reference to a SymbolGenerator object, which assigns a number to
     * a rule to allow for quick comparison
     */
    public SymbolGenerator symGen;

    /**
     * A reference to the head finder rules.
     */
    public HeadFinderRules rules;
        
    /**
     * Creates a new Grammar object, and sets the StartSymbol and EndSymbol to
     * "S!" and "/S" respectively
     *
     * @param pg The production generator to use
     * @param sg The symbol generator to use
     */
    public Grammar(ProductionGenerator pg, SymbolGenerator sg) {
        this.prodGen = pg;
        this.symGen = sg;
        rules = new HeadFinderRules(symGen);
        StartSymbol = symGen.create("S!");
        EndSymbol = symGen.create("/S");
        lexicon.put(sg.create("/S"), sg.create("/S"));
    }
    
    /**
     * Adds a chain to the grammar.
     *
     * @param c The <code>Chain</code> to add.
     */
    public void addChain(Chain c) {
        // If there aren't any chains who have c's tail already in the grammar,
        // add it to the Hashtable of chains
        if(!chains.containsKey(c.getTail())) {
            chains.put(c.getTail(), new Hashtable<Symbol, HashSet<Chain>>());
        // If there aren't any chains who have c's head already in the grammar,
        // add it to the Hashtable of chains
        } if(!chains.get(c.getTail()).containsKey(c.getHead())) {
            chains.get(c.getTail()).put(c.getHead(), new HashSet<Chain>());
        // If the chain c hasn't yet been seen, put it in the HashSet, and set the
        // times seen to 1
        } if(!chains.get(c.getTail()).get(c.getHead()).contains(c)) {
            chains.get(c.getTail()).get(c.getHead()).add(c);
            chainCounts.put(c, new Integer(1));
        // Increment the number of times teh chain has been seen
        } else {
            chainCounts.put(c, new Integer(chainCounts.get(c).intValue() + 1));
        }
    }
    
    /**
     * Add the allowable triple <parent, leftSibling, attachmentNode> to the list
     * of allowable triples.
     */
    public void addAllowableTriple(Symbol parent, Symbol leftSibling, Symbol attachmentNode) {
        Tuple context = new Tuple(2);
        context.put(0, parent);
        context.put(1, leftSibling);
        if(!allowableTriples.containsKey(attachmentNode))
            allowableTriples.put(attachmentNode, new HashSet<Tuple>());
        allowableTriples.get(attachmentNode).add(context);
    }
    
    /**
     * @return a String containing the XML representation of the grammar in this
     * grammar object
     */
    public String getXMLRep() {
        String ret = "<?xml version=\"1.0\"?>\n<grammar type=\"perceptron\">\n";
        ret += parameters.getXMLRep();
        for(Enumeration<Symbol> en1 = chains.keys(); en1.hasMoreElements();) {
            Symbol t = en1.nextElement();
            for(Enumeration<Symbol> en2 = chains.get(t).keys(); en2.hasMoreElements();) {
                Symbol h = en2.nextElement();
                for(Iterator<Chain> it = chains.get(t).get(h).iterator(); it.hasNext();) {
                    ret += it.next().getXMLRep();
                }
            }
        }
        for(Enumeration<Symbol> en = allowableTriples.keys(); en.hasMoreElements();) {
            Symbol t3 = en.nextElement();
            for(Iterator<Tuple> it = allowableTriples.get(t3).iterator(); it.hasNext();) {
                Tuple tup = it.next();
                Symbol t1 = (Symbol)tup.get(0);
                Symbol t2 = (Symbol)tup.get(1);
                ret += "<allowable parent=\"" + t1 + "\" sibling=\"" + t2 + "\" node=\"" + t3 + "\"/>\n";
            }
        }
//        ret += "</productions>\n";
        ret += "</grammar>\n";
        return ret;
    }

    /**
     * Write the XML representation of the grammar to a File.
     */
    public void writeXML(File file) {
        FileWriter fw;
        try {
            fw = new FileWriter(file);
            fw.write("<?xml version=\"1.0\"?>\n<grammar type=\"perceptron\">\n");
            fw.write(parameters.getXMLRep());
            for(Enumeration<Symbol> en1 = chains.keys(); en1.hasMoreElements();) {
                Symbol t = en1.nextElement();
                for(Enumeration<Symbol> en2 = chains.get(t).keys(); en2.hasMoreElements();) {
                    Symbol h = en2.nextElement();
                    for(Iterator<Chain> it = chains.get(t).get(h).iterator(); it.hasNext();) {
                        fw.write(it.next().getXMLRep());
                    }
                }
            }
            for(Enumeration<Symbol> en = allowableTriples.keys(); en.hasMoreElements();) {
                Symbol t3 = en.nextElement();
                for(Iterator<Tuple> it = allowableTriples.get(t3).iterator(); it.hasNext();) {
                    Tuple tup = it.next();
                    Symbol t1 = (Symbol)tup.get(0);
                    Symbol t2 = (Symbol)tup.get(1);
                    if(t1.toString().equals("&"))
                        t1 = symGen.create("&amp;");
                    if(t2.toString().equals("&"))
                        t2 = symGen.create("&amp;");
                    if(t3.toString().equals("&"))
                        t3 = symGen.create("&amp;");
                    fw.write("<allowable parent=\"" + t1 + "\" sibling=\"" + t2 + "\" node=\"" + t3 + "\"/>\n");
                }
            }
//            fw.write("</productions>\n");
            fw.write("</grammar>\n");
            fw.close();
        } catch (IOException ex) {
            ex.printStackTrace();
        }
    }

    /**
     * Method used for debugging purposes; prints out all chains in the grammar
     */
    public void printChains() {
        for(Enumeration<Symbol> out = chains.keys(); out.hasMoreElements(); ) {
            Symbol tail = out.nextElement();
            Hashtable<Symbol, HashSet<Chain>> inner = chains.get(tail);
            for(Enumeration<Symbol> in = inner.keys(); in.hasMoreElements(); ) {
                Symbol head = in.nextElement();
                HashSet<Chain> cs = inner.get(head);
                for(Iterator<Chain> it = cs.iterator(); it.hasNext();) {
                    System.out.println("TAIL: " + tail + " HEAD: " + head + " CHAIN: " + it.next());
                }
            }
        }        
    }
    
    /**
     * @return a String containing the representation of the grammar in this
     * grammar object
     */
    public String toString() {
        String ret = "-------------------------------------------\n";
        ret += parameters;
        ret += "-------------------------------------------\n";

//        for(Iterator<Production> it = productions.iterator(); it.hasNext();) {
//            ret += it.next();
//        }
        ret += "-------------------------------------------\n";
        for(Enumeration<Symbol> out = chains.keys(); out.hasMoreElements(); ) {
            Symbol tail = out.nextElement();
            Hashtable<Symbol, HashSet<Chain>> inner = chains.get(tail);
            for(Enumeration<Symbol> in = inner.keys(); in.hasMoreElements(); ) {
                Symbol head = in.nextElement();
                HashSet<Chain> cs = inner.get(head);
                for(Iterator<Chain> it = cs.iterator(); it.hasNext();) {
                    Chain c = it.next();
                    System.err.println(c.getHead() + " == " + head + ":  " + c.getHead().equals(head));
                    System.err.println(c.getTail() + " == " + tail + ":  " + c.getTail().equals(tail));
                    if(!c.getHead().equals(head) || !c.getTail().equals(tail))
                        ret += "!!! ";
                    ret += "(" + head + ", " + tail + ") " + c;
                }
            }
        }
        ret += "-------------------------------------------\n";
        for(Enumeration<Symbol> out = allowableTriples.keys(); out.hasMoreElements();) {
            Symbol trip3 = out.nextElement();
            for(Iterator<Tuple> it = allowableTriples.get(trip3).iterator(); it.hasNext();) {
                Tuple trip = it.next();
                Symbol trip1 = (Symbol)trip.get(0);
                Symbol trip2 = (Symbol)trip.get(1);
                ret += "<" + trip1 + "," + trip2 + "," + trip3 + ">\n";
            }
        }
        ret += "-------------------------------------------\n";
        return ret;
    }
    
    /**
     * Adds a parameter to the parameter vector
     *
     * @param f the feature to add
     * @param i the value to add
     */
    public void addParameter(Feature f, double i) {
        this.parameters.addFeature(f, i);
    }
    
    /**
     * Retrieve all the chains indexed by the head Symbol and tail Symbol, and convert
     * each chain into a Node tree.
     *
     * @param head The head of the chain to search for
     * @param tail The tail of the chain to seasrch for
     * @param lookAhead The look ahead token to use for firing features when building the node trees.
     * @param lookAheadPOS The POS tag of the look ahead token to use for firing features when building the node trees.
     */
    public LinkedList<Node> getChains(Symbol head, Symbol tail, Symbol lookAhead, Symbol lookAheadPOS) {
        // Get all the chains with the specified head and tail.
        HashSet<Chain> returnchains = chains.get(tail).get(head);
        LinkedList<Node> ret = new LinkedList<Node>();
        // Iterate through all the chains found, and build the corresponding node trees.
        for(Iterator<Chain> it = returnchains.iterator(); it.hasNext();) {
            Chain next = it.next();
            ret.addLast(next.buildNodeTree(parameters, symGen, lookAhead, lookAheadPOS, rules));
        }
        return ret;
    }
    
    /**
     * Retrieve all the chains indexed by the tail Symbol, and convert
     * each chain into a Node tree.
     *
     * @param tail The tail of the chain to seasrch for
     * @param lookAhead The look ahead token to use for firing features when building the node trees.
     * @param lookAheadPOS The POS tag of the look ahead token to use for firing features when building the node trees.
     */
    public LinkedList<Node> getChains(Symbol tail, Symbol lookAhead, Symbol lookAheadPOS) {
        // Get the Hashtable of HashSets containing all the chains indexed by tail
        Hashtable<Symbol, HashSet<Chain>> returnchains = chains.get(tail);
        LinkedList<Node> ret = new LinkedList<Node>();
        // Go through all the possible heads of chains
        for(Enumeration<Symbol> next = returnchains.keys(); next.hasMoreElements();) {
            // Retrieve all the chains indexed by the head and tail, and add them
            // to the return LinkedList
            LinkedList<Node> r = getChains(next.nextElement(), tail, lookAhead, lookAheadPOS);
            ret.addAll(r);
        }
        return ret;
    }
    
    /**
     * Add a production rule to the list of productions in the grammar.
     *
     * TODO: When adding the Terminals/NonTerminals to the set, implement it so
     * that it would keep track of the number of times a terminal or non-terminal
     * has been seen.
     */
    public void addProduction(Rule r) {
        // Add the production to the list
        productions.add(prodGen.create(r));
        // Add possible symbols to the sets of terminals and non-terminals.
        if(r.hasTerminalRHS()) {
            lexicon.put(r.RHS().get(0), r.LHS());
            Terminals.add(r.RHS().get(0));
        } else {
            for(Iterator<Symbol> it = r.RHS().iterator(); it.hasNext();) {
                nonTerminals.add(it.next());
            }
        }
        nonTerminals.add(r.LHS());
    }
    
    /**
     * Adds a list of chains to the grammar
     *
     * @param chains a LinkedList of chains to add
     */
    public void addChains(LinkedList<Chain> chains) {
        for(Iterator<Chain> it = chains.iterator(); it.hasNext();) {
            this.addChain(it.next());
        }
    }
    
    /**
     * This method adds a rule to the grammar, and should derive all the possible
     * chains and allowable triples. 
     *
     * @param r the Rule to add to the grammar.
     */
    // Note that this doesn't quite work. It needs
    // to be fixed so that the chain formation is not dependent on the order in
    // which the rules are read. Also needs to properly deal with left recursion
    // and unit recursion in rules. I wrote a method to do this, but it was broken,
    // so I removed it.
    public void addRule(Rule r) {
        // add the rule to the set of productions
        productions.add(prodGen.create(r));
        if(r.hasTerminalRHS()) {
            lexicon.put(r.RHS().get(0), r.LHS());
        }
        // add the parent of the production to the list of nonterminals
        nonTerminals.add(r.LHS());
        // If the rule is a lexical item, then add the children to the list of
        // terminal symbols
        if(r.hasTerminalRHS()) {
            for(Iterator<Symbol> i = r.RHS().iterator(); i.hasNext();)
                Terminals.add(i.next());
        } else {
            // add the first child to the list of nonterminals
            Symbol i = r.RHS().get(0);
            nonTerminals.add(i);
            // iterate through all the children of the production
            for(int c = 1; c < r.RHS().size(); c++) {
                Symbol j = r.RHS().get(c);
                // add the child to the set of nonterminals
                nonTerminals.add(j);
                if(!allowableTriples.containsKey(j))
                    allowableTriples.put(j, new HashSet<Tuple>());
                Tuple constraint = new Tuple(2);
                constraint.put(0, r.LHS());
                constraint.put(1, i);
                allowableTriples.get(j).add(constraint);
                i = j;
            }
        }
        Chain chain = new Chain(r.LHS());
        chain.append(r.RHS().get(0));
        // If there is no chain stored whose tail is the head of the new chain,
        // then add the new chain to the list of chains
        if(!chains.containsKey(chain.getHead())) {
            this.addChain(chain);
        } else {
            Hashtable<Symbol, HashSet<Chain>> expChains = chains.remove(chain.getHead());
            for(Enumeration<Symbol> en = expChains.keys(); en.hasMoreElements();) {
                HashSet<Chain> chainList = expChains.get(en.nextElement());
                for(Iterator<Chain> it = chainList.iterator(); it.hasNext();) {
                    Chain curChain = it.next();
                    curChain.append((Chain)chain.subList(1, chain.size()));
                    this.addChain(curChain);
                }
            }
        }
    }
    
}
